"""
灵芯派截图工具

功能说明：
1. 支持通过三击触摸事件触发截图功能（需接入触碰传感器）。
2. 支持手动截图和查看已有截图。
3. 使用 pygame 显示截图，适配灵芯派屏幕尺寸（480x800）。

模块依赖：
- smartpi: 用于读取触碰传感器数据（可选）。
- PIL (ImageGrab): 用于屏幕截图。
- pygame: 用于显示截图。
- os/sys: 文件操作和系统功能。

函数说明：
- take_screenshot(): 截图并保存到指定目录。
- show_screenshot(image_path): 使用 pygame 显示指定路径的截图。
- detect_clicks(): 监听三击事件并触发截图（需接入触碰传感器）。
- main(): 主程序入口，提供模式选择。

注意事项：
1. 截图默认保存在 ./screenshots 目录下。
2. 显示截图时需确保 pygame 和依赖库正确安装。
3. 若使用触碰传感器，需调整 CLICK_THRESHOLD 参数。
"""

from smartpi import touch_sensor
import time
from PIL import ImageGrab
import os
import pygame
import sys

# 截图保存目录
SCREENSHOT_DIR = "screenshots"

# 创建截图目录
if not os.path.exists(SCREENSHOT_DIR):
    os.makedirs(SCREENSHOT_DIR)

# 点击检测参数
CLICK_THRESHOLD = 50  # 触摸数据阈值（根据调试结果调整）
TRIPLE_CLICK_INTERVAL = 0.5  # 三击间隔（秒）

# 记录点击事件
click_times = []

def take_screenshot():
    """截图并保存到指定目录"""
    timestamp = time.strftime("%Y%m%d_%H%M%S")
    screenshot_path = os.path.join(SCREENSHOT_DIR, f"screenshot_{timestamp}.png")
    
    # 使用 ffmpeg 捕获屏幕
    try:
        import subprocess
        subprocess.run([
            "ffmpeg",
            "-f", "x11grab",
            "-video_size", "480x800",
            "-i", ":0.0",
            "-vframes", "1",
            "-q:v", "2",
            screenshot_path
        ], check=True)
        print(f"截图已保存: {screenshot_path}")
        return screenshot_path
    except Exception as e:
        print(f"截图失败: {e}")
        return None

def show_screenshot(image_path):
    """使用 pygame 显示截图"""
    try:
        print("正在初始化 pygame...")
        pygame.init()
        
        # 设置窗口尺寸为灵芯派屏幕尺寸（480x800）
        print("正在设置窗口尺寸...")
        screen = pygame.display.set_mode((480, 800))
        pygame.display.set_caption("截图预览")
        
        # 加载图片并缩放到窗口尺寸
        print(f"正在加载图片: {image_path}")
        try:
            image = pygame.image.load(image_path)
            print("图片加载成功，正在缩放...")
            image = pygame.transform.scale(image, (480, 800))
        except Exception as e:
            print(f"图片加载失败: {e}")
            pygame.quit()
            return
        
        print("开始显示图片...")
        running = True
        while running:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                elif event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_ESCAPE:
                        running = False
            
            # 填充背景色（避免黑屏）
            screen.fill((0, 0, 0))
            screen.blit(image, (0, 0))
            pygame.display.flip()
            
            # 检查窗口是否被关闭
            if not pygame.display.get_init():
                running = False
        
        pygame.quit()
        print("图片显示结束")
    except Exception as e:
        print(f"无法显示图片: {e}")

def detect_clicks():
    """检测三击事件并截图"""
    global click_times
    
    print("启动三击截图监听程序...")
    print("按下 Ctrl+C 停止程序")
    
    try:
        while True:
            touch_data = touch_sensor.get_value(1)
            
            if touch_data > CLICK_THRESHOLD:
                click_times.append(time.time())
                print("检测到触摸事件")
                time.sleep(0.1)  # 防抖
            
            # 检测三击
            if len(click_times) >= 3:
                time_diff1 = click_times[-1] - click_times[-2]
                time_diff2 = click_times[-2] - click_times[-3]
                
                if time_diff1 <= TRIPLE_CLICK_INTERVAL and time_diff2 <= TRIPLE_CLICK_INTERVAL:
                    print("检测到三击，正在截图...")
                    screenshot_path = take_screenshot()
                    print(f"是否查看截图？ (y/n)")
                    choice = input().strip().lower()
                    if choice == 'y':
                        show_screenshot(screenshot_path)
                    click_times = []
                else:
                    click_times = []  # 重置
            
            time.sleep(0.05)
    except KeyboardInterrupt:
        print("程序已停止")

def main():
    """主程序入口"""
    print("请选择模式:")
    print("1. 监听三击截图")
    print("2. 查看已有截图")
    choice = input("输入选项 (1/2): ").strip()
    
    if choice == "1":
        # detect_clicks()
        take_screenshot()
    elif choice == "2":
        if not os.listdir(SCREENSHOT_DIR):
            print("截图目录为空，请先截图。")
            return
        
        print("可用的截图:")
        screenshots = sorted(os.listdir(SCREENSHOT_DIR))
        for idx, file in enumerate(screenshots, 1):
            print(f"{idx}. {file}")
        
        selected = input("输入要查看的截图编号: ").strip()
        try:
            selected_idx = int(selected) - 1
            if 0 <= selected_idx < len(screenshots):
                show_screenshot(os.path.join(SCREENSHOT_DIR, screenshots[selected_idx]))
            else:
                print("无效的编号。")
        except ValueError:
            print("请输入有效数字。")
    else:
        print("无效选项。")

if __name__ == "__main__":
    main()